import os
import time
import cv2
import numpy as np
import pandas as pd
from keras.applications.vgg19 import VGG19, preprocess_input
from keras.preprocessing import image
from keras.models import Model
from keras.layers import Dense, Flatten, BatchNormalization, Dropout, GlobalAveragePooling2D
from sklearn.metrics import classification_report, roc_curve, auc
import matplotlib.pyplot as plt
from keras.utils import to_categorical
from keras.optimizers import Adam
from keras.losses import categorical_crossentropy
from keras.callbacks import EarlyStopping, ModelCheckpoint, ReduceLROnPlateau
from keras.preprocessing.image import ImageDataGenerator

# Load VGG19 model
base_model = VGG19(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
x = base_model.output
x = GlobalAveragePooling2D()(x)
x = BatchNormalization()(x)
x = Dense(512, activation='relu', kernel_regularizer='l2')(x)
x = Dropout(0.5)(x)
x = Dense(256, activation='relu', kernel_regularizer='l2')(x)
x = Dropout(0.5)(x)
output = Dense(5, activation='softmax')(x)  # new output layer for 5 classes
model = Model(inputs=base_model.input, outputs=output)

# Freeze the layers except the last few layers
for layer in base_model.layers[:-4]:
    layer.trainable = False

# Compile the model
model.compile(optimizer=Adam(learning_rate=0.0001), loss=categorical_crossentropy, metrics=['accuracy'])

# Print the model summary
model.summary()

# Function to detect and crop chest area
def detect_and_crop_chest(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    chest_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_fullbody.xml')
    chests = chest_cascade.detectMultiScale(gray, 1.1, 4)
    for (x, y, w, h) in chests:
        return image[y:y+h, x:x+w]
    return image  # return original image if no chest is detected

# Load and preprocess images
def load_images_from_folder(folder):
    images = []
    for filename in os.listdir(folder):
        img = cv2.imread(os.path.join(folder, filename))
        if img is not None:
            img = detect_and_crop_chest(img)
            img = cv2.resize(img, (224, 224))
            images.append(img)
    return images

# Load training images
train_dir = 'D:\\kaggle\\final\\class\\train'
class_dirs = ['adenocarcinoma', 'squamous', 'big', 'small', 'normal']

train_images = []
train_labels = []

for i, class_dir in enumerate(class_dirs):
    class_images = load_images_from_folder(os.path.join(train_dir, class_dir))
    train_images.extend(class_images)
    train_labels.extend([i] * len(class_images))

# Load testing images
test_dir = 'D:\\kaggle\\final\\class\\test'

test_images = []
test_labels = []

for i, class_dir in enumerate(class_dirs):
    class_images = load_images_from_folder(os.path.join(test_dir, class_dir))
    test_images.extend(class_images)
    test_labels.extend([i] * len(class_images))

# Combine data
X_train = np.array(train_images)
y_train = np.array(train_labels)

X_test = np.array(test_images)
y_test = np.array(test_labels)

# Preprocess images
X_train = preprocess_input(X_train)
X_test = preprocess_input(X_test)

# Convert labels to categorical
y_train = to_categorical(y_train, num_classes=5)
y_test = to_categorical(y_test, num_classes=5)

# Data augmentation
datagen = ImageDataGenerator(
    rotation_range=30,
    width_shift_range=0.3,
    height_shift_range=0.3,
    shear_range=0.3,
    zoom_range=0.3,
    horizontal_flip=True,
    fill_mode='nearest'
)

# Callbacks
checkpoint_dir = '../checkpoints'
os.makedirs(checkpoint_dir, exist_ok=True)

checkpoint = ModelCheckpoint(filepath=os.path.join(checkpoint_dir, 'chest_CT_SCAN-VGG19.h5'),
                             monitor='val_loss',
                             mode='auto',
                             save_best_only=True)
early_stopping = EarlyStopping(verbose=1, patience=10)
reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.2, patience=5, min_lr=0.00001)

# Train model
start_time = time.time()
history = model.fit(datagen.flow(X_train, y_train, batch_size=32), epochs=100, validation_data=(X_test, y_test),
                    callbacks=[checkpoint, early_stopping, reduce_lr])
end_time = time.time()

# Predict on test set
y_pred = model.predict(X_test)

# Convert predictions to labels
y_pred_labels = np.argmax(y_pred, axis=1)
y_test_labels = np.argmax(y_test, axis=1)

# Print classification report
print(classification_report(y_test_labels, y_pred_labels))

# Calculate and plot ROC curve and AUC for each class
for i in range(5):
    fpr, tpr, _ = roc_curve(y_test[:, i], y_pred[:, i])
    roc_auc = auc(fpr, tpr)
    plt.figure()
    plt.plot(fpr, tpr, label='ROC curve (area = %0.2f)' % roc_auc)
    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'Receiver Operating Characteristic for Class {i}')
    plt.legend(loc="lower right")
    plt.show()

# Print program running time
print(f'Program running time: {end_time - start_time}')

'''
              precision    recall  f1-score   support

           0       0.36      0.70      0.47        30
           1       0.00      0.00      0.00        12
           2       0.41      0.23      0.29        31
           3       0.71      0.67      0.69        36
           4       1.00      0.77      0.87        30

    accuracy                           0.54       139
   macro avg       0.49      0.47      0.46       139
weighted avg       0.57      0.54      0.53       139

              precision    recall  f1-score   support

           0       0.41      0.50      0.45        30
           1       0.07      0.08      0.08        12
           2       0.42      0.32      0.36        31
           3       0.59      0.67      0.62        36
           4       1.00      0.77      0.87        30

    accuracy                           0.53       139
   macro avg       0.50      0.47      0.48       139
weighted avg       0.55      0.53      0.53       139

'''